{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e7436800",
   "metadata": {},
   "source": [
    "# VM instrument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1bc1eee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "\n",
    "from tvm import relax\n",
    "from tvm.relax.testing import nn\n",
    "from tvm.relax.testing.lib_comparator import LibCompareVMInstrument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "34c0b1ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_exec(data_shape):\n",
    "    builder = relax.BlockBuilder()\n",
    "    weight1_np = np.random.randn(64, 64).astype(\"float32\")\n",
    "    weight2_np = np.random.randn(64, 64).astype(\"float32\")\n",
    "\n",
    "    with builder.function(\"main\"):\n",
    "        model = nn.Sequential(\n",
    "            nn.Linear(data_shape[1], weight1_np.shape[0], bias=False),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(weight2_np.shape[0], weight2_np.shape[1], bias=False),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        data = nn.Placeholder(data_shape, name=\"data\")\n",
    "        output = model(data)\n",
    "        params = [data] + model.parameters()\n",
    "        builder.emit_func_output(output, params=params)\n",
    "\n",
    "    mod = builder.get()\n",
    "\n",
    "    params = {\"linear_weight\": weight1_np, \"linear_weight1\": weight2_np}\n",
    "    mod = relax.transform.BindParams(\"main\", params)(mod)\n",
    "\n",
    "    target = \"llvm\"\n",
    "    return tvm.compile(mod, target)\n",
    "\n",
    "\n",
    "def get_exec_int32(data_shape):\n",
    "    builder = relax.BlockBuilder()\n",
    "\n",
    "    with builder.function(\"main\"):\n",
    "        model = nn.ReLU()\n",
    "        data = nn.Placeholder(data_shape, dtype=\"int32\", name=\"data\")\n",
    "        output = model(data)\n",
    "        params = [data] + model.parameters()\n",
    "        builder.emit_func_output(output, params=params)\n",
    "\n",
    "    mod = builder.get()\n",
    "    target = \"llvm\"\n",
    "    return tvm.compile(mod, target)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a397392",
   "metadata": {},
   "source": [
    "## 测试 conv2d_cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e3ee34f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_4023119/2478082253.py:17: UserWarning: Returning type `vm.Storage` which is not registered via register_object, fallback to Object\n",
      "  vm[\"main\"](tvm.nd.array(data_np))\n"
     ]
    }
   ],
   "source": [
    "data_np = np.random.randn(1, 64).astype(\"float32\")\n",
    "ex = get_exec(data_np.shape)\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "hit_count = {}\n",
    "\n",
    "def instrument(func, name, before_run, ret_val, *args):\n",
    "    if (name, before_run) not in hit_count:\n",
    "        hit_count[(name, before_run)] = 0\n",
    "    hit_count[(name, before_run)] += 1\n",
    "    assert callable(func)\n",
    "    if before_run:\n",
    "        assert ret_val is None\n",
    "    if name == \"matmul\":\n",
    "        return relax.VMInstrumentReturnKind.SKIP_RUN\n",
    "\n",
    "vm.set_instrument(instrument)\n",
    "vm[\"main\"](tvm.nd.array(data_np))\n",
    "assert hit_count[(\"matmul\", True)] == 2\n",
    "assert (\"matmul\", False) not in hit_count\n",
    "assert hit_count[(\"relu\", True)] == 2\n",
    "assert hit_count[(\"relu\", False)] == 2\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0032c34",
   "metadata": {},
   "source": [
    "## 测试 {class}`~tvm.relax.testing.lib_comparator.LibCompareVMInstrument`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ef956ea6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_4023119/2973990986.py:7: UserWarning: Returning type `vm.Storage` which is not registered via register_object, fallback to Object\n",
      "  vm[\"main\"](tvm.nd.array(data_np))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tvm.nd.NDArray shape=(1, 64), cpu:0>\n",
       "array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 1, 0, 0, 0, 0, 0, 1,\n",
       "        1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 2,\n",
       "        0, 0, 0, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],\n",
       "      dtype=int32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_np = np.random.randn(1, 64).astype(\"int32\")\n",
    "ex = get_exec_int32(data_np.shape)\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "# compare against library module\n",
    "cmp = LibCompareVMInstrument(vm.module.imported_modules[0], tvm.cpu(), verbose=False)\n",
    "vm.set_instrument(cmp)\n",
    "vm[\"main\"](tvm.nd.array(data_np))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4573fc3b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
